/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once
#include "beanmachine/graph/distribution/distribution.h"

namespace beanmachine {
namespace distribution {

/*
  A distribution representing the mixture of two distributions D1 and D2
  with probability P1 of sampling from D1 and 1 - P1 of sampling from D2.

  It must have exactly three parents.
  The first parent for bimixture distribution must be the probability P1.
  The second and third parents must be distributions D1 and D2.
  The sample type must be consistent with the distribution parents.
*/
class Bimixture : public Distribution {
 public:
  Bimixture(
      graph::ValueType sample_type,
      const std::vector<graph::Node*>& in_nodes);
  Bimixture(
      graph::AtomicType sample_type,
      const std::vector<graph::Node*>& in_nodes);
  ~Bimixture() override {}
  bool _bool_sampler(std::mt19937& gen) const override;
  double _double_sampler(std::mt19937& gen) const override;
  graph::natural_t _natural_sampler(std::mt19937& gen) const override;

  double log_prob(const graph::NodeValue& value) const override;
  void log_prob_iid(const graph::NodeValue& value, Eigen::MatrixXd& log_probs)
      const override;
  void gradient_log_prob_value(
      const graph::NodeValue& value,
      double& grad1,
      double& grad2) const override;
  void gradient_log_prob_param(
      const graph::NodeValue& value,
      double& grad1,
      double& grad2) const override;

  void backward_value(
      const graph::NodeValue& value,
      graph::DoubleMatrix& back_grad,
      double adjunct = 1.0) const override;
  void backward_value_iid(
      const graph::NodeValue& value,
      graph::DoubleMatrix& back_grad) const override;
  void backward_value_iid(
      const graph::NodeValue& value,
      graph::DoubleMatrix& back_grad,
      Eigen::MatrixXd& adjunct) const override;

  void backward_param(const graph::NodeValue& value, double adjunct = 1.0)
      const override;
  void backward_param_iid(const graph::NodeValue& value) const override;
  void backward_param_iid(
      const graph::NodeValue& value,
      Eigen::MatrixXd& adjunct) const override;
};

} // namespace distribution
} // namespace beanmachine
